#include "stdafx.h"
#include "LuaASM.h"

typedef unsigned char BYTE;
typedef BYTE *PBYTE;
typedef signed int          INT32;
typedef unsigned int        UINT32;
typedef unsigned long		DWORD;
inline PBYTE DetourGenPush(PBYTE pbCode, UINT32 nValue)
{
	*pbCode++ = 0x68;
	*((UINT32*&)pbCode)++ = nValue;
	return pbCode;
}

inline PBYTE DetourGenPushad(PBYTE pbCode)
{
	*pbCode++ = 0x60;
	return pbCode;
}

inline PBYTE DetourGenPopad(PBYTE pbCode)
{
	*pbCode++ = 0x61;
	return pbCode;
}
inline PBYTE DetourGenCall(PBYTE pbCode, PBYTE pbJmpDst, PBYTE pbJmpSrc = 0)
{
	if (pbJmpSrc == 0) {
		pbJmpSrc = pbCode;
	}
	*pbCode++ = 0xE8;
	*((INT32*&)pbCode)++ = pbJmpDst - (pbJmpSrc + 5);
	return pbCode;
}

inline PBYTE DetourGenBreak(PBYTE pbCode)
{
	*pbCode++ = 0xcc;
	return pbCode;
}

inline PBYTE DetourGenRet(PBYTE pbCode)
{
	*pbCode++ = 0xc3;
	return pbCode;
}  
inline PBYTE DetourGenAddEspByte(PBYTE pbCode,BYTE value)
{
	*pbCode++ = 0x83;
	*pbCode++ = 0xc4;
	*pbCode++ = value;
	return pbCode;
}
//////////////////////////////////////////////////////////////////////////
struct ASMCode{
	char Code[500];
	char Source[500];
	PBYTE CodeWritePosition;
	PBYTE SrcWritePosition;
};
static int luaNewAsmCode(lua_State* L){
	ASMCode* pCode = (ASMCode*)lua_newuserdata(L,sizeof(ASMCode));
	pCode->CodeWritePosition = (PBYTE)pCode->Code;
	pCode->SrcWritePosition = (PBYTE)pCode->Code;

	luaL_getmetatable(L,"AsmCode.CodeMeta");
	lua_setmetatable(L,-2);
	return 1;
}
static int luaAsmReset(lua_State* L){
	ASMCode* pCode = (ASMCode*)lua_newuserdata(L,sizeof(ASMCode));
	pCode->CodeWritePosition = (PBYTE)pCode->Code;
	pCode->SrcWritePosition =  (PBYTE)pCode->Code;
	return 0;
}

static ASMCode* luaL_GetAsmCodeObject( lua_State* L ) 
{
	ASMCode* pCode = (ASMCode*)luaL_checkudata(L,1,"AsmCode.CodeMeta");
	luaL_argcheck(L, pCode != NULL, 1, "`AsmCode' expected"); 
	return pCode;
}

static int luaAsmPush(lua_State* L){
	ASMCode* pCode = luaL_GetAsmCodeObject(L);

	int argType = lua_type(L,2);
	switch(argType){
	case LUA_TNUMBER:
		{
			unsigned int arg = luaL_checkint(L,2);
			pCode->CodeWritePosition = DetourGenPush((PBYTE)pCode->CodeWritePosition,arg);
			break;
		}
	case LUA_TSTRING:
		{
			const char* arg=luaL_checkstring(L,2);
			pCode->CodeWritePosition = DetourGenPush((PBYTE)pCode->CodeWritePosition,(DWORD)arg);
			break;
		}
	};
	
	return 0;
}

static int luaAsmPushad(lua_State* L){
	ASMCode* pCode = luaL_GetAsmCodeObject(L);
	
	pCode->CodeWritePosition = DetourGenPushad((PBYTE)pCode->CodeWritePosition);
	return 1;
}
static int luaAsmPopad(lua_State* L){
	ASMCode* pCode = luaL_GetAsmCodeObject(L);

	pCode->CodeWritePosition = DetourGenPopad((PBYTE)pCode->CodeWritePosition);
	return 1;
}
static int luaAsmGenCall(lua_State* L){
	ASMCode* pCode = luaL_GetAsmCodeObject(L);

	int arg = luaL_checkint(L,2);

	pCode->CodeWritePosition = DetourGenCall((PBYTE)pCode->CodeWritePosition,(PBYTE)arg);
	return 0;
}
static int luaAsmAddEsp(lua_State* L){
	ASMCode* pCode = luaL_GetAsmCodeObject(L);
	
	int arg = luaL_checkint(L,2);
	pCode->CodeWritePosition = DetourGenAddEspByte(pCode->CodeWritePosition,arg);
	return 0;
}
static int luaAsmDo(lua_State* L){
	ASMCode* pCode = luaL_GetAsmCodeObject(L);

	typedef DWORD (*TargetFunc_type)();
	TargetFunc_type TargetFunc = (TargetFunc_type)((PBYTE)pCode->Code);
	
	DWORD oldFlag;
	VirtualProtect((void*)TargetFunc,1000,PAGE_EXECUTE_READWRITE,&oldFlag);
	DWORD r = TargetFunc();
	VirtualProtect((void*)TargetFunc,1000,oldFlag,&oldFlag);
	lua_pushnumber(L,r);
	return 1;
}
static int luaAsmRet(lua_State* L){
	ASMCode* pCode = luaL_GetAsmCodeObject(L);

	pCode->CodeWritePosition = DetourGenRet(pCode->CodeWritePosition);
	return 0;
}
//////////////////////////////////////////////////////////////////////////
static const struct luaL_reg AsmCodeLib [] = { 
	{"New",luaNewAsmCode},
	{"Reset",luaAsmReset},
	{"Push",luaAsmPush},
	{"Pushad",luaAsmPushad},
	{"Popad",luaAsmPopad},
	{"Call",luaAsmGenCall},
	{"AddEsp",luaAsmAddEsp},
	{"Ret",luaAsmRet},	

	{"Do",luaAsmDo},
	{NULL, NULL} 
};
bool OpenLuaASM(lua_State* lua )
{
	
	luaI_openlib(lua,"AsmCode",AsmCodeLib,0);
	luaL_newmetatable(lua,"AsmCode.CodeMeta");
	return true;
}
